//
//  Exploit.m
//  Dopamine
//
//  Created by Lars Fröder on 10.01.24.
//

#import "DOExploit.h"
#import <dlfcn.h>
#import <sys/utsname.h>
#import <sys/sysctl.h>
#import <UIKit/UIKit.h>
#import <libjailbreak/machine_info.h>

#import "DOExploitManager.h"
#import "DOEnvironmentManager.h"

@implementation DOExploit

- (instancetype)initWithPath:(NSString *)exploitPath info:(NSDictionary *)exploitInfo flavorName:(NSString *)flavorName flavorDictionary:(NSDictionary *)flavorInfo
{
    self = [super init];
    if (self) {
        _flavorInfo = flavorInfo;
        _flavorName = flavorName;
        _executablePath = [exploitPath stringByAppendingPathComponent:exploitInfo[@"CFBundleExecutable"]];
        
        NSString *typeString = exploitInfo[@"DPExploitType"];
        if ([typeString isEqualToString:@"Kernel"]) {
            _type = EXPLOIT_TYPE_KERNEL;
        }
        else if ([typeString isEqualToString:@"PAC"]) {
            _type = EXPLOIT_TYPE_PAC;
        }
        else if ([typeString isEqualToString:@"PPL"]) {
            _type = EXPLOIT_TYPE_PPL;
        }
        else {
            NSLog(@"Unknown Exploit Type: %@", typeString);
            return nil;
        }
        
        NSNumber *priorityNum = flavorInfo[@"DPFlavorPriority"];
        _priority = [priorityNum unsignedLongLongValue];
        
        _name = exploitInfo[@"CFBundleDisplayName"];
        _identifier = exploitInfo[@"CFBundleIdentifier"];
        if (![flavorName isEqualToString:@"default"]) {
            _name = [NSString stringWithFormat:@"%@, %@", _name, flavorName];
            _identifier = [NSString stringWithFormat:@"%@.%@", _identifier, flavorName];
        }
    }
    return self;
}

- (NSString *)displayName
{
    if (self.flavorName) {
        return [NSString stringWithFormat:@"%@ (%@)", self.name, self.flavorName];
    }
    else {
        return self.name;
    }
}

- (BOOL)isSupported
{
    if (!_supportLoaded) {
        struct utsname systemInfo;
        uname(&systemInfo);
        NSString *thisDevice = [NSString stringWithUTF8String:systemInfo.machine];
        
        cpu_subtype_t cpuFamily = 0;
        size_t cpuFamilySize = sizeof(cpuFamily);
        sysctlbyname("hw.cpufamily", &cpuFamily, &cpuFamilySize, NULL, 0);
        
        NSString *thisCPU;
        switch (cpuFamily) {
            case CPUFAMILY_ARM_TYPHOON:
                thisCPU = @"A8";
                break;
            case CPUFAMILY_ARM_TWISTER:
                thisCPU = @"A9";
                break;
            case CPUFAMILY_ARM_HURRICANE:
                thisCPU = @"A10";
                break;
            case CPUFAMILY_ARM_MONSOON_MISTRAL:
                thisCPU = @"A11";
                break;
            case CPUFAMILY_ARM_VORTEX_TEMPEST:
                thisCPU = @"A12";
                break;
            case CPUFAMILY_ARM_LIGHTNING_THUNDER:
                thisCPU = @"A13";
                break;
            case CPUFAMILY_ARM_FIRESTORM_ICESTORM:
                thisCPU = @"A14"; // Also M1
                break;
            case CPUFAMILY_ARM_BLIZZARD_AVALANCHE:
                thisCPU = @"A15"; // Also M2
                break;
            case CPUFAMILY_ARM_EVEREST_SAWTOOTH:
                thisCPU = @"A16";
                break;
            case CPUFAMILY_ARM_COLL:
                thisCPU = @"A17";
                break;
        }
        
        char OSVersionString[64];
        size_t OSVersionStringLen = sizeof(OSVersionString) - 1;
        sysctlbyname("kern.osversion", OSVersionString, &OSVersionStringLen, NULL, 0);
        NSString *thisIOSBuild = [NSString stringWithUTF8String:OSVersionString];
        
        NSString *thisIOSVersion = [[UIDevice currentDevice] systemVersion];
        BOOL isArm64e = [[DOEnvironmentManager sharedManager] isArm64e];
        
        NSArray *supportedRanges = _flavorInfo[@"DPSupportedRanges"];
        for (NSDictionary *supportedRange in supportedRanges) {
            NSString *rangeStart = supportedRange[@"Start"];
            NSString *rangeEnd = supportedRange[@"End"];
            
            NSComparisonResult startResult = [thisIOSVersion compare:rangeStart options:NSNumericSearch];
            NSComparisonResult endResult = [thisIOSVersion compare:rangeEnd options:NSNumericSearch];
            
            _supported = ((endResult == NSOrderedAscending || endResult == NSOrderedSame) && (startResult == NSOrderedDescending || startResult == NSOrderedSame));
            if (_supported) break;
        }
        
        NSArray *exclusionInclusionSets = _supported ? _flavorInfo[@"DPSupportExclude"] : _flavorInfo[@"DPSupportInclude"];
        if (exclusionInclusionSets) {
            for (NSDictionary *set in exclusionInclusionSets) {
                NSArray *devicesSet = set[@"Devices"];
                BOOL affectsThisDevice = NO;
                if (devicesSet) {
                    for (NSString *device in devicesSet) {
                        if (([device isEqualToString:@"arm64e"] && isArm64e) ||
                            ([device isEqualToString:@"arm64"] && !isArm64e) ||
                            [device isEqualToString:thisDevice] ||
                            [device isEqualToString:thisCPU]) {
                            affectsThisDevice = YES;
                            break;
                        }
                    }
                }
                else {
                    // No "Devices" means all
                    affectsThisDevice = YES;
                }
                if (!affectsThisDevice) continue;
                
                NSArray *builds = set[@"Builds"];
                if (builds) {
                    for (NSString *build in builds) {
                        if ([thisIOSBuild isEqualToString:build]) {
                            _supported = !_supported;
                            break;
                        }
                    }
                }
                else {
                    _supported = !_supported;
                }
                break;
            }
        }
        _supportLoaded = YES;
    }
    return _supported;
}

- (int)load
{
    if (!_handle) {
        _handle = dlopen(_executablePath.fileSystemRepresentation, RTLD_NOW);
    }
    if (_handle) return 0;
    return -1;
}

- (int)run
{
    if (!_handle) return -1;
    int (*exploit_init)(const char *flavor) = dlsym(_handle, "exploit_init");
    if (!exploit_init) return -1;
    int r = exploit_init(_flavorName.UTF8String);
    if (r == 0) [[DOExploitManager sharedManager].activeExploits addObject:self];
    return r;
}

- (int)cleanup
{
    if (!_handle) return -1;
    int (*exploit_deinit)(void) = dlsym(_handle, "exploit_deinit");
    if (!exploit_deinit) return -1;
    int r = exploit_deinit();
    if (r == 0) [[DOExploitManager sharedManager].activeExploits removeObject:self];
    return r;
}

- (NSString *)description
{
    return [NSString stringWithFormat:@"<Exploit: identifier = %@>", self.identifier];
}

@end
